#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File     : loader.py
@Project  : pipecoco
@Date     : 2021/8/13
@Author   : Zhang Jinyang
@Contact  : zhang-jy@sjtu.edu.cn
'''

from pipecoco.pipecoco_layer import *

# 用于识别融合网络特征和激活方式的结构列表
merge_layer = ['Flatten','Dense']
activate_layer = ['ReLU','BatchNorm2d','Pad']
fused_layer = ['Conv2d','MaxPool2d']


def layers_partitioner(layers_list):
    """
    把layers_list转换成可融合的网络部分和不可融合的网络部分
    """

    layer_partition = []
    fused_layers = []
    non_fused_layers = []
    i = 0
    while i < len(layers_list):
        layer = layers_list[i]
        if isinstance(layer, str):
            # 为了识别bypass结构，在列表中用[..., bypass, activate, h, ...]
            # 表示将当前层和前h层输出相加后执行activate计算的bypass操作
            if layer == 'bypass':
                fused_layers.append(Bypass(layers_list[i + 1], layers_list[i + 2], 3))
                i += 2
        elif isinstance(layer, list):
            # 还有部分不属于mindspore中nn类的算子结构，要进行封装后才能以网络层模式执行计算
            non_fused_layers.append(Reducemean(layer[0], layer[1]))

        elif layer.cls_name in merge_layer:
            # 到达融合网络终点
            non_fused_layers.extend(layers_list[i:])
            break
        else:
            # 将融合网络中的主要操作和之后的激活、正则化等操作识别为融合网络的结构单元。
            block = [layer]

            while i + 1 < len(layers_list):
                if hasattr(layers_list[i + 1], 'cls_name') and layers_list[i + 1].cls_name in activate_layer:
                    block.append(layers_list[i + 1])
                    i += 1
                elif (hasattr(layers_list[i + 1], 'name') and layers_list[i + 1].name in activate_layer):
                    block.append(getattr(nn, layers_list[i + 1].name)())
                    i += 1
                else:
                    break
            fused_layers.append(block)
        i += 1
    layer_partition.append(fused_layers)
    layer_partition.append(nn.SequentialCell(non_fused_layers))
    return layer_partition


def dfs_search_layer(net, layer_list):
    """
    用深度搜索的方式找到所有基本网络结构，并加入layer_list列表中
    """

    if len(net.cells())==0:
        # 没有子结构时加入列表
        layer_list.append(net)
    else:
        for cell in net.cells():
            if isinstance(cell, nn.SequentialCell):
                # 序列执行的网络结构，依次搜索列表中各网络
                for layer in cell:
                    dfs_search_layer(layer, layer_list)
            else:
                dfs_search_layer(cell, layer_list)

